"""
仿照matplotlib的概念，用pygame实现了一点绘图功能(Scope)：
    1. 单变量序列的折线图
    2. 双变量的散点图、x-y图形（将相邻两散点连接起来）

关于坐标系：
    计算机图形的坐标原点都是左上角，x轴正方向向右，y轴正方向向下。为了
    方便计算，所有坐标的计算都假设原点在左下角，最后再将整幅图像关于x轴
    对称翻转过来。不便处理的是文字的方向，需要在每一处用到文字的地方都
    翻转一次。

Stone at GDOU, Sep. 2020
"""
import time
import sys
from collections import deque

import pygame as pyg



pyg.init()
BGC = (0,0,0)  # background color


class BaseAxis:
    def __init__(self, sfc, lim_lower=0, lim_upper=2.5):
        """:params sfc: a pygame surface object"""
        self.sfc = sfc
        self._sizex, self._sizey = self.sfc.get_size()

        self.bgc = BGC
        self.color = (0, 255, 255)
        self.font_color = (255, 255, 0)
        
        self._font_size = 26
        self.font = pyg.font.SysFont('arial', self.font_size)
        self._font_height = self.font.get_height()

        self._lim_lower = lim_lower
        self._lim_upper = lim_upper
        
        self._tick_step = 1         # how many value per tick
        self._tick_thickness = 2    # pixels
        self.tick_length = 6       # pixels, the drawing result is affected by sfc's size
        
        self.spine_padding = 30     # pixels, the distance between spine and the border of ``sfc``, where to place the tick label
        self.spine_thickness = 2    # pixels
        
        self.draw()
    
    def calculate_ticks_to_draw(self):
        """Calculate the informations needed to draw the ticks. This algorithm will
        give prioty to draw the tick of value 0, that is, if zero fall in the range of ``lim_lower``
        and ``lim_upper`` , then make sure tick of zero will be draw. 
        
        Return the tick values that will be draw."""
        ticks = []  # tick values to draw
        
        start = self._lim_upper - (self._lim_upper % self._tick_step)  # start from where to zero could be devided in integer
        
        n = int((self._lim_upper - self._lim_lower) / self._tick_step)  # how many ticks could the range be divided into
        for i in range(n+1):
            tick = start - i*self._tick_step
            if tick >= self._lim_lower:
                ticks.append(tick)
        
        return ticks        
    
    def val2coor(self, val):
        """Map the val to coordinate in pixels of the corresponding axis."""
        self.raise_not_implemented_error()
    
    def draw_tick(self):
        self.raise_not_implemented_error()
    
    def draw_spine(self):
        self.raise_not_implemented_error()
    
    def draw(self):
        self.draw_spine()
        self.draw_tick()
    
    def raise_not_implemented_error(self):
        raise NotImplementedError("This drawing method should be implemented in the specific child class")
    
    @property
    def font_size(self):
        return self._font_size
    
    @font_size.setter
    def font_size(self, val):
        self._font_size = val
        self.font = pyg.font.SysFont('arial', self.font_size)
        self._font_height = self.font.get_height()
    
    @property
    def lim_upper(self):
        return self._lim_upper

    @lim_upper.setter
    def lim_upper(self, val):
        self._lim_upper = val
        self.sfc.fill(self.bgc)
        self.draw_spine()
        self.draw_tick()

    @property
    def lim_lower(self):
        return self._lim_lower
    
    @lim_lower.setter
    def lim_lower(self, val):
        self._lim_lower = val
        self.sfc.fill(self.bgc)
        self.draw_spine()
        self.draw_tick()

    @property
    def tick_step(self):
        return self._tick_step

    @tick_step.setter
    def tick_step(self, val):
        self._tick_step = val
        self.draw()


class YAxis(BaseAxis):
    def __init__(self, sfc, lim_lower=0, lim_upper=2.5):
        super(YAxis, self).__init__(sfc, lim_lower, lim_upper)
        # self.font_offset_x = 5  # the offset relative to tick where to draw the font
        # self.font_offset_y = 2
        self.spine_padding = 50
    
    def val2coor(self, val):
        """Map the val to coordinate in pixels of the corresponding axis."""
        return int(self._sizey * (val - self._lim_lower) / (self._lim_upper - self._lim_lower))    
        
    def draw_tick(self):
        ticks = self.calculate_ticks_to_draw()
        range_ = self._lim_upper - self._lim_lower
        x1 = self.spine_padding + self.spine_thickness
        x2 = x1 + self.tick_length
        
        for tick in ticks:
            y = (tick - self._lim_lower) / range_ * self._sizey
            pyg.draw.line(self.sfc, self.color, (x1, y), (x2, y), self._tick_thickness)
            
            # draw text
            txt = '%.2f' % tick
            txt = self.font.render(txt, False,  self.font_color)
            txt = pyg.transform.flip(txt, False, True)
            self.sfc.blit(txt, (0, y))
        
    def draw_spine(self):
        pyg.draw.line(self.sfc, self.color, (self.spine_padding,0), (self.spine_padding, self._sizey), 2)


class XAxis(BaseAxis):
    def __init__(self, sfc, lim_lower=0, lim_upper=2.5):
        super(XAxis, self).__init__(sfc, lim_lower, lim_upper)
        
        self.spine_padding = 30
    
    def val2coor(self, val):
        """Map the val to coordinate in pixels of the corresponding axis."""
        return int(self._sizex * (val - self._lim_lower) / (self._lim_upper - self._lim_lower))
        
    def draw_tick(self):
        ticks = self.calculate_ticks_to_draw()
        range_ = self._lim_upper - self._lim_lower
        y1 = self.spine_padding + self.spine_thickness
        y2 = y1 + self.tick_length
        
        for tick in ticks:
            x = (tick - self._lim_lower) / range_ * self._sizex
            pyg.draw.line(self.sfc, self.color, (x, y1), (x, y2), self._tick_thickness)
            
            # draw text
            txt = '%.2f' % tick
            txt = self.font.render(txt, False,  self.font_color)
            txt = pyg.transform.flip(txt, False, True)
            self.sfc.blit(txt, (x, 0))
        
    def draw_spine(self):
        pyg.draw.line(self.sfc, self.color, (0, self.spine_padding), (self._sizex, self.spine_padding), 2)


class BaseCurve:
    """A ``Curve`` object is used to store some parameters and data
    that is necessary to draw the diagram. This base class defines
    some interfaces that the ``Curve`` classes have to implement.
    """
    def __init__(self, axes, color):
        self.axes = axes
        self.color = color
        self.thickness = 1  # the thickness of curve
    
    # def append_sample(self, val):
        # raise NotImplementedError("This method is not implemented by child class.")
    
    def append_sample(self, val):
        self.axes.append_sample(val, curve=self)


class CurveYT(BaseCurve):
    def __init__(self, axes, color):
        super(CurveYT, self).__init__(axes, color)
        self.last_pos = None


class CurveXY(BaseCurve):
    def __init__(self, axes, color):
        super(CurveXY, self).__init__(axes, color)
        self.n_points = 80      # store how many points of sample
        self._points = deque()
        
        self.connected = False # whethet draw lines connecting the dots. If not, it becomes scatter diagram
        self.draw_point = True # whethet draw the dots.
        
        self.pt_radius = 2      # radius of point
    
    def append_point(self, point):
        self._points.append(point)
        if len(self._points) > self.n_points:
            self._points.popleft()
    
    @property
    def points(self):
        return self._points


class AxesDraw:
    """
    Using strategy pattern.
    Encapsule the functions of drawing a diagram in order to keep the clarity.
    """
    def __init__(self, sfc, axes):
        self.sfc = sfc
        self.axes = axes
        self._sizex, self._sizey = sfc.get_size()
        
        # parameters in y-t mode
        self.x_step = 1  # used in y-t mode
    
    def draw_yt(self, val, curve):
        """draw the curve in y-t mode, the curve object must be ``CurveYT`` type."""
        if not isinstance(curve, CurveYT):
            raise TypeError("Except a <CurveYT>, but got a {}".format(type(curve)))
        
        # calculate coordinates
        y = self.axes.yaxis.val2coor(val)
        if curve.last_pos is None:
            x = 0
            curve.last_pos = (x, y)  # initial value
        else:
            x = curve.last_pos[0] + self.x_step
        
        # boundary check
        if x > self._sizex:
            x = 0
            curve.last_pos = (x,y)
            self._clear()
        
        # do drawing
        pyg.draw.line(self.sfc, curve.color, curve.last_pos, (x, y), curve.thickness)
        
        # update
        curve.last_pos = (x, y)
    
    def draw_xy(self, pair, curve):
        """Draw the ``pair`` of value in x-y mode"""
        if not isinstance(curve, CurveXY):
            raise TypeError("Except a <CurveXY>, but got a {}".format(type(curve)))
        
        # calculate coordinates
        x = self.axes.xaxis.val2coor(pair[0])
        y = self.axes.yaxis.val2coor(pair[1])
        
        # update
        curve.append_point((x, y))

        # do drawing
        self.sfc.fill(self.axes.bgc)
        if curve.connected:
            points = curve.points
            for i in range(len(points) - 1):
                pyg.draw.line(self.sfc, curve.color, points[i], points[i+1], curve.thickness)
        
        if curve.draw_point:
            for p in curve.points:
                pyg.draw.circle(self.sfc, curve.color, p, curve.pt_radius)
        
    def _clear(self):
        self.sfc.fill(self.axes.bgc)


class Axes:
    """Composition of instances like the axis, curves and so on."""
    def __init__(self, sfc):
        self.sfc = sfc
        
        self.mode_xy = False
        self.bgc = BGC
        self.color = (255, 0, 255, 255)
        sizex, sizey = sfc.get_size()
        self._sizex = sizex
        self._sizey = sizey
        self._width_yaxis = 60             # pixels
        self._height_xaxis = 50            # pixels
        self._ogn = (self._width_yaxis, self._height_xaxis)  # pixels, coordinates origin corresponding to drawing
        
        # y axis 和 x axis的绘图区（surface）都贴附于axes.sfc的边上，因此subsurface的起点坐标
        # 中，有一个坐标为0。而另一个坐标为``self._ogn``的一个坐标，这是为了使坐标轴错开，以便
        # 绘制坐标轴交点处的刻度。数值转换成坐标的计算由各axis负责，然后再由axes加上``self._ogn``并绘制出来。
        # 因此 x axis 的左端点和 y axis 的右端点一般不重合，但是坐标与数值的对应关系是对的。
        # 
        # 此外，``draw_sfc``用于绘制曲线，给其subsurface的原点就在self._ogn,以免去在坐标计算公式中
        # 引入原点偏移的麻烦。
        yleft, ytop, ywidth, yheight = 0,self._ogn[1], self._width_yaxis, sizey-self._ogn[1]
        xleft, xtop, xwidth, xheight = self._ogn[0],0, sizex-self._ogn[0], self._height_xaxis
        dleft, dtop, dwidth, dheight = self._ogn[0], self._ogn[1], sizex - self._ogn[0], sizey - self._ogn[1]
        
        y_axis_sfc = self.sfc.subsurface([yleft, ytop, ywidth, yheight])
        x_axis_sfc = self.sfc.subsurface([xleft, xtop, xwidth, xheight])
        drawer_sfc = self.sfc.subsurface([dleft, dtop, dwidth, dheight])
        self.yaxis = YAxis(y_axis_sfc)
        self.xaxis = XAxis(x_axis_sfc)
        self.drawer = AxesDraw(drawer_sfc, self)
        
        self.default_curve = CurveYT(self, self.color)
        self.curves = [self.default_curve]
    
    def append_sample(self, val, curve=None):
        """Draw the new sample on the figure with proper drawing mode."""
        if curve is None:
            curve = self.curves[0]
        
        if isinstance(curve, CurveYT):
            self.drawer.draw_yt(val, curve)
        elif isinstance(curve, CurveXY):
            self.drawer.draw_xy(val, curve)
        else:
            raise TypeError("Unkonw type of curve: {}".format(type(curve)))

    def create_curve(self, color, curve_type='y-t'):
        """Crete a new curve and register it in a list. Tough the mechanism does not
        prohibit drawing curve in y-t mode and x-y mode at the same time, and it even will work,
        it is not recommanded to do so because it might make people confused and cause 
        problems.
        :param curve_type: pass in string 'y-t' or 'x-y'.
        """
        if curve_type == 'y-t':
            curve = CurveYT(self, color)
        elif curve_type == 'x-y':
            curve = CurveXY(self, color)
        else:
            raise ValueError("Unknow curve type:{}. Available types are: 'y-t', 'x-y'")
        
        self.curves.append(curve)

        return curve
    
    def redraw(self):
        """clear and redraw the image"""
        self.sfc.fill(self.bgc)
        self.yaxis.draw()
        
        if self.mode_xy:
            self.xaxis.draw()


class TrackBar:
    def __init__(self, sfc, min_value=20, max_value=400):
        self._value = min_value  # the value of TrackBar
        self.min_value = min_value
        self.max_value = max_value

        # size and width and height
        self.sizex = sfc.get_width()
        self.sizey = sfc.get_height()

        self.text_area_w = self.sizex
        self.text_area_h = 30
        self.motion_area_w = self.sizex
        self.motion_area_h = self.sizey-self.text_area_h
        self.bar_w = self.sizex
        self.bar_h = 10
        self.rail_w = 6
        self.rail_h = self.motion_area_h
        self.slide_h = self.rail_h - self.bar_h  # the range where the bar could slide

        # colors
        self.font_color = (255, 255, 255)
        self.bar_color = (228, 44, 100)
        self.rail_color = (200, 200, 200)
        self.bgc = (0, 200, 0)

        # surfaces to draw
        self.sfc = sfc
        self.text_area = self.sfc.subsurface([0, 0, self.text_area_w, self.text_area_h])
        self.motion_area = self.sfc.subsurface([0, self.text_area_h, self.motion_area_w, self.motion_area_h])

        # some elements of widget
        self.bar = pyg.Surface([self.bar_w, self.bar_h])
        self.rail = pyg.Surface([self.rail_w, self.rail_h])
        self.bar.fill(self.bar_color)
        self.rail.fill(self.rail_color)

        # font
        self.font = pyg.font.SysFont('arial', 24)
        
        self.fresh()

    def event_handler(self, events):
        for event in events:
            # only MOUSEMOTION has parameter ``buttons``, so this order in if statement is significant,
            # with the fearture of lasy logic.
            if event.type == pyg.MOUSEMOTION \
               and event.buttons[0] == 1:

                x, y = event.pos
                ox, oy = self.motion_area.get_abs_offset()
                delta_y = y - oy
                value = self.min_value + (self.slide_h - delta_y)/self.slide_h * (self.max_value - self.min_value)

                if value > self.max_value:
                    self.value = self.max_value
                elif value < self.min_value:
                    self.value = self.min_value
                else:
                    self.value = value

    def fresh(self):
        """Fresh the appearance"""
        self.sfc.fill(self.bgc)

        # draw elements
        x, y = int(self.sizex/2 - self.rail_w/2), 0
        self.motion_area.blit(self.rail, (x, y))

        x, y = 0, int((self._value - self.min_value)/(self.max_value - self.min_value) * self.slide_h)
        y = self.slide_h - y  # flilp vertically
        self.motion_area.blit(self.bar, (x,y))

        # draw text
        txt = self.font.render('%i' % self.value, False, self.font_color)
        txt_w = txt.get_width()
        x = int(self.text_area_w - txt_w) / 2
        self.text_area.blit(txt, [x, 0])

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        self._value = value
        self.fresh()


class Scope:
    def __init__(self, sizex=600, sizey=300):
        self._sizex = sizex
        self._sizey = sizey
        
        self.scr = pyg.display.set_mode([sizex, sizey])
        self._event_handlers = []

        panel_width, panel_height = 46, 220
        panel_sfc = self.scr.subsurface([sizex - panel_width, 20, panel_width, panel_height])
        self.panel = TrackBar(panel_sfc)
        self._register_event_handler(self.panel.event_handler)
        
        axes_sfc = pyg.Surface([sizex - panel_width, sizey])
        self.axes = Axes(axes_sfc)

    def _handle_event(self):
        """Run in main loop"""
        events = pyg.event.get()
        for event in events:
            if event.type == pyg.QUIT\
               or (event.type==pyg.KEYDOWN and event.key == pyg.K_ESCAPE):
                print('exit')
                sys.exit()

        # throw the events to the handlers and let them handle it
        # TODO this is inefficient, as every handler has to filter the unwanted events
        for func in self._event_handlers:
            func(events)

    def _register_event_handler(self, func):
        self._event_handlers.append(func)
        
    def show(self):
        sfc = pyg.transform.flip(self.axes.sfc, False, True)
        self.scr.blit(sfc, [0, 0])
        
        pyg.display.update()

    def append_sample(self, val):
        self.axes.append_sample(val)
    
    def apply(self):
        """Make some appearance settings via attributes come to effect."""
        self.axes.redraw()
        self.panel.fresh()
        
    def mainloop(self, user_task, *args, **kw):
        print("start main loop")
        while True:
            self._handle_event()
            user_task()


def test_loop():
    cos_curve = scope.axes.create_curve([255, 255, 0])
    cxy = scope.axes.create_curve((255,255,0), 'x-y')
    cxy.connected = True
    # cxy.draw_point = False

    while True:
        for i in range(1, 942):
            pyg.time.delay(10)
            y = math.sin(i/15) +1
            cos_y = 0.5*math.cos(i/5) +1
            scope.axes.append_sample(y)
            cos_curve.append_sample(cos_y)
            # scope.panel.value = y
            
            cxy.append_sample((cos_y, y))
            scope.show()

            yield


if __name__ == '__main__':
    pyg.display.set_caption('六一儿童节快乐')
    scope = Scope()
    
    scope.axes.yaxis.tick_step = 1
    scope.axes.yaxis.lim_lower = -1
    scope.axes.yaxis.lim_upper = 5.1
    scope.axes.mode_xy = True
    
    scope.panel.min_value = -1
    scope.panel.max_value = 3
    scope.apply()
    
    
    import math
    scope.mainloop((lambda it:lambda: next(it))(test_loop()) )


